package com.csust.zerocode.core.handler;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.csust.zerocode.ai.model.message.*;
import com.csust.zerocode.ai.tools.BaseTool;
import com.csust.zerocode.ai.tools.ToolManager;
import com.csust.zerocode.constant.AppConstant;
import com.csust.zerocode.core.builder.VueProjectBuilder;
import com.csust.zerocode.model.entity.User;
import com.csust.zerocode.model.enums.ChatHistoryMessageTypeEnum;
import com.csust.zerocode.service.ChatHistoryService;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.HashSet;
import java.util.Set;

/**
 * JSON 消息流处理器
 * 处理 VUE_PROJECT 类型的复杂流式响应，包含工具调用信息
 */
@Slf4j
@Component
public class JsonMessageStreamHandler implements StreamHandler {

    @Resource
    private ToolManager toolManager;
    /**
     * 处理 TokenStream（VUE_PROJECT）
     * 解析 JSON 消息并重组为完整的响应格式
     *
     * @param originFlux         原始流
     * @param chatHistoryService 聊天历史服务
     * @param appId              应用ID
     * @param loginUser          登录用户
     * @return 处理后的流
     */
    public Flux<String> handle(Flux<String> originFlux,
                               ChatHistoryService chatHistoryService,
                               long appId, User loginUser) {
        // 收集数据用于生成后端记忆格式
        StringBuilder aiChatHistoryStringBuilder = new StringBuilder();
        StringBuilder  frontChatHistoryStringBuilder = new StringBuilder();
        // 用于跟踪已经见过的工具ID，判断是否是第一次调用
        Set<String> seenToolIds = new HashSet<>();
        return originFlux
                .map(chunk -> {
                    // 解析每个 JSON 消息块 为AiResponse
                    AiResponse aiResponse = JSONUtil.toBean(handleJsonMessageChunk(chunk, seenToolIds), AiResponse.class);
                    //得到消息类型
                    StreamMessageTypeEnum streamMessageTypeEnum = StreamMessageTypeEnum.getEnumByValue(aiResponse.getType());
                    String content = aiResponse.content;
                    // 收集完整信息方便前端进行显示
                    frontChatHistoryStringBuilder.append(content);
                    if(StreamMessageTypeEnum.AI_RESPONSE.equals(streamMessageTypeEnum)){
                        //如果是ai响应
                        aiChatHistoryStringBuilder.append(content);
                    }
                    return content;
                })
                .filter(StrUtil::isNotEmpty) // 过滤空字串
                .doOnComplete(() -> {
                    // 流式响应完成后，添加 AI 消息到对话历史
                    String aiResponse = aiChatHistoryStringBuilder.toString();
                    String frontResponse = frontChatHistoryStringBuilder.toString();
                    if(StrUtil.isNotBlank(aiResponse)) {
                        //添加 AI消息
                        chatHistoryService.addChatMessage(appId, aiResponse, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    }
                    if(StrUtil.isNotBlank(frontResponse)) {
                        // 添加 前端展示消息
                        chatHistoryService.addChatMessage(appId, frontResponse, ChatHistoryMessageTypeEnum.FRONTEND.getValue(), loginUser.getId());
                    }
                    // 异步构造 Vue 项目
//                    String projectPath = AppConstant.CODE_OUTPUT_ROOT_DIR + "/vue_project_" + appId;
//                    vueProjectBuilder.buildProjectAsync(projectPath);
                })
                .doOnError(error -> {
                    // 如果AI回复失败，也要记录错误消息
                    String errorMessage = "AI回复失败: " + error.getMessage();
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.FRONTEND.getValue(), loginUser.getId());
                });
    }

    /**
     * 解析并收集 TokenStream 数据
     * @param chunk 数据块
     * @param seenToolIds 用于记录已经见过的工具ID
     * @return 处理后的数据块
     */
    private String handleJsonMessageChunk(String chunk, Set<String> seenToolIds) {
        // 解析 JSON，先还原为基类  StreamMessage
        StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
        //继而根据基类的 type字段获取具体消息类型
        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
        switch (typeEnum) {
            case AI_RESPONSE -> {
                AiResponseMessage aiMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
                String data = aiMessage.getData();
                // 直接拼接响应
//                chatHistoryStringBuilder.append(data);
                return getAiMsgJson(data, typeEnum);
            }
            case TOOL_REQUEST -> {
                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(chunk, ToolRequestMessage.class);
                String toolName = toolRequestMessage.getName();
                String toolId = toolRequestMessage.getId();
                // 检查是否是第一次看到这个工具 ID
                if (toolId != null && !seenToolIds.contains(toolId)) {
                    // 第一次调用这个工具，记录 ID 并完整返回工具信息
                    seenToolIds.add(toolId);
                    // 根据工具名称获取工具实例
                    BaseTool tool = toolManager.getTool(toolName);
                    // 返回格式化的工具调用信息
                    return getAiMsgJson(tool.generateToolRequestResponse(), typeEnum);
                } else {
                    // 不是第一次调用这个工具，直接返回空
                    return getAiMsgJson("", typeEnum);
                }
            }
            case TOOL_EXECUTED -> {
                ToolExecutedMessage toolExecutedMessage = JSONUtil.toBean(chunk, ToolExecutedMessage.class);
                JSONObject jsonObject = JSONUtil.parseObj(toolExecutedMessage.getArguments());
                String toolName = toolExecutedMessage.getName();
                // 根据工具名称获取工具实例
                BaseTool tool = toolManager.getTool(toolName);
                // 返回格式化的工具调用信息
                String result = tool.generateToolExecutedResult(jsonObject);
                // 输出前端和要持久化的内容
                String output = String.format("\n\n%s\n\n", result);
//                chatHistoryStringBuilder.append(output);
                return getAiMsgJson(output, typeEnum);
            }
            default -> {
                log.error("不支持的消息类型: {}", typeEnum);
                return "";
            }
        }
    }

    /**
     * 得到 {AiResponse} 的JSON格式
     * @param content 消息内容
     * @param typeEnum 消息类型
     * @return JSON格式的消息
     */
    private String getAiMsgJson(String content, StreamMessageTypeEnum typeEnum) {
        AiResponse aiResponse = AiResponse.builder().content(content).type(typeEnum.getValue()).build();
        return JSONUtil.toJsonStr(aiResponse);
    }

    // AiResponse 解决Ai 学习消息历史（包含工具调用）从而 “偷懒” 假装调用了工具
    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class AiResponse {
        private String content;

        /**
         * @see StreamMessageTypeEnum
         */
        private String type;
    }

}
